#include "epoll.h"
#include "LockAndUnlock.h"

void CEpollServer::StartUp()
{
  TraceLog("begin start")
  // 1. sock 三部曲
  _listenfd = socket(AF_INET, SOCK_STREAM, 0);
  if(-1 == _listenfd){
    ErrorLog("Create listenfd");
    return;
  }

  struct sockaddr_in local;
  local.sin_family = AF_INET;
  local.sin_port = htons(_port);
  local.sin_addr.s_addr = htonl(INADDR_ANY);

  if(-1 == bind(_listenfd, (struct sockaddr*)&local, sizeof(struct sockaddr))){
    ErrorLog("Bind Socket");
    return;
  }

  if(-1 == listen(_listenfd, 99999)){
    ErrorLog("Listen Socket");
    return;
  }

  // 2. 创建 epoll 事件表
  _eventfd = epoll_create(_MAX_EVENT);
  if(-1 == _eventfd){
    ErrorLog("Create Event");
    return;
  }
  
  // 3. 将 listen 的读事件添加到事件表中
  OperateEvent(_listenfd, EPOLLIN, EPOLL_CTL_ADD);

  // 4. 循环等待事件
  EventLoop();
}

void CEpollServer::EventLoop()
{
  TraceLog("Begin Event Loop");

  struct epoll_event events[_MAX_EVENT];
  for(; ;){
    int size = epoll_wait(_eventfd, events, _MAX_EVENT, -1);
    if(-1 == size){
      ErrorLog("epoll_wait");
    }

    for(int i = 0; i < size; ++i){
      int fd = events[i].data.fd;
      // 1. 监听套接字读就绪：有新的链接请求
      if(_listenfd == fd){
        struct sockaddr addr;
        socklen_t length;
        int client_fd = accept(_listenfd, &addr, &length);
        if(-1 == client_fd){
          ErrorLog("accept client");
        }
        TraceLog("client connect");
        _DoConnectEventHandle(client_fd);
      }
      // 2. 读事件就绪：有数据来了
      else if(events[i].events & EPOLLIN){
        _DoReadEventHandle(events[i].data.fd);
      }
      // 3. 写事件就绪：对应的接收端可以接受数据了
      else if(events[i].events & EPOLLOUT){
        _DoWriteEventHandle(events[i].data.fd);
      }
      else{
        ErrorLog("Events Error!!!");
      }
    }
  } // end for(;;)
}


void CEpollServer::SendInLoop (int fd, const char* buf, size_t length)
{
  int send_length = send(fd, buf, length, 0);
  // 1. 发送失败
  if(-1 == send_length){
    ErrorLog("Send to %d", fd);
    return;
  }
  // 2. 一次未能发送完，添加写事件
  else if(send_length < length){
    Connect* connect = _connectMap[fd];
    Channel* channel = &connect->_clientChannel;
    // 判断是往哪端写
    if(connect->_serverChannel._fd == fd){
      channel = &connect->_serverChannel;
    }
    channel->_buffer.append(buf+send_length);
    channel->_event |= EPOLLOUT;
    channel->_event |= EPOLLONESHOT;
    OperateEvent(fd, channel->_event, EPOLL_CTL_MOD);
  }
}

void CEpollServer::_DoWriteEventHandle(int fd)
{
  // 1. 找到要往哪个连接里面写
  std::map<int, Connect*>::iterator it = _connectMap.find(fd);
  if(it != _connectMap.end()){
    Connect* connect = it->second;
    Channel* channel = &connect->_serverChannel;
    if(fd == connect->_clientChannel._fd){
      channel = &connect->_clientChannel;
    }
    // 2. 不能直接使用连接中的buffer
    std::string tmp;
    tmp.swap(channel->_buffer);
    SendInLoop(fd, tmp.c_str(), tmp.size());
  }
  // 3. 走到此处表明连接出错了
  else{
    assert(false);
  }
}

void CEpollServer::Forward(Channel* client_channel, Channel* server_channel, bool recv_decrypt, bool send_encrypt)
{
  const int length = 5120;
  char buf[length];
  size_t recv_length = recv(client_channel->_fd, buf, length, 0);
  // 1. 读取数据个数小于零，recv 出错
  if(recv_length < 0){
    ErrorLog("recv from client_channel");
  }
  // 2. 收到 0 个数据，表示发送方已断开连接,"他"已经不会给我发数据了
  //    "我"不会再收到上家数据,也就不会往下家写,往下家写的端口就可以关了
  else if(0 == recv_length){
    shutdown(server_channel->_fd, SHUT_WR);
    RemoveConnect(client_channel->_fd);
  }
  // 3. 正常收到数据,往下家转发
  else{
    // 3.1 对收到的数据进行加密, 再转发出去
    //     如: transfer 收到来自浏览器的数据, 先进行加密再发给 socks5 服务器
    if(recv_decrypt){
      Decrypt(buf, recv_length);
    }
    // 3.2 对要发送的数据进行解密，再转发
    //     如: socks5 服务器收到来自 transfer 数据, 进行解密后再发给 google
    if(send_encrypt){
      Encrypt(buf, recv_length);
    }

    // 4. 开始循环发送
    buf[recv_length] = '\0';
    SendInLoop(server_channel->_fd, buf, recv_length);
  }
}

void CEpollServer::RemoveConnect(int fd)
{
  TraceLog("delete %d", fd);
  // 1. 找到要删除的连接
  std::map<int, Connect*>::iterator it = _connectMap.find(fd);
  if(it != _connectMap.end()){
    Connect* connect = it->second;
    Channel* channel = &connect->_clientChannel;
    if(connect->_serverChannel._fd == fd){
      channel = &connect->_serverChannel;
    }
    // 2. 删除 event 事件表中相关事件
    //    并将连接从 map 中删除
    OperateEvent(fd, channel->_event, EPOLL_CTL_DEL);
    _connectMap.erase(it);

    // 3. 释放连接所占用的空间(每个连接是 new 出来的)
    if(0 == --connect->_counter){
      delete connect;
    }
  }
}
